// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
//

#include <kllm/kai/tokenize.h>

namespace kllm {

    void KaiTokenize::tokenize(const KaiRequest &req, KaiResponse&res) {
        CHECK(req.query_type() == QUERY_TOKENIZE);
        res.set_model_alias(_context->params.model_alias);
        auto add_special = req.tokenize_request().add_special();
        auto with_pieces = req.tokenize_request().with_pieces();
        turbo::Status status;
        std::vector<std::vector<llama_token>> tokens = _context->tokenize_input_prompts(req.tokenize_request().prompts(), add_special, with_pieces, status);
        if(!status.ok()) {
            res.mutable_status()->set_errmsg(std::string(status.message()));
            res.mutable_status()->set_code(status.raw_code());
        }
        for(auto &l1 :tokens ) {
            for(auto &l2 : l1) {
                res.mutable_tokens()->Add(l2);
            }
        }
        res.set_with_pieces(with_pieces);
        res.set_add_special(add_special);
        res.set_query_type(QUERY_TOKENIZE);
    }

    void KaiTokenize::detokenize(const KaiRequest &req, KaiResponse&res) {
        CHECK(req.query_type() == QUERY_DETOKENIZE);
       auto content = _context->tokens_to_str(req.detokenize().tokens().cbegin(), req.detokenize().tokens().cend());
        res.set_content(content);
        res.mutable_status()->set_code(0);
        res.mutable_status()->set_errmsg("ok");
    }
}